import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset


class TransformerDataset(Dataset):
    def __init__(self, source_df:pd.DataFrame, target_df:pd.DataFrame, target_idx:int, seq:int, device):
        super(TransformerDataset, self).__init__()
        self.source = source_df.to_numpy(np.float32)
        self.source = torch.tensor(self.source, dtype=torch.float32, device=device)
        self.target = target_df.to_numpy(np.float32)
        self.target = torch.tensor(self.target, dtype=torch.float32, device=device)
        self.seq = seq
        self.target_idx = target_idx
        self.features = self.source.shape[-1]
        self.device = device

    def __getitem__(self, index):
        src = self.source[index:index + self.seq]
        tgt = torch.concat((torch.zeros([1, self.features]).to(self.device), self.target[index:index + self.seq - 1]), 0)
        tgt_y = self.target[index:index + self.seq, self.target_idx]
        return src, tgt, tgt_y

    def __len__(self):
        return len(self.source) - self.seq